import os
import cv2
import random
import numpy as np
from PIL import Image, ImageEnhance, ImageOps, ImageFile
from function.viewImage import *
# from function.readTxtFile import *
# from function.readTxtFile import readFileList, XYWH_2_XYXY
from reader.Reader_landmark_label import getMirrorLM, getMirrorBox

PI = 3.14159265358

def lm_image_enhance(image, landmark):
    '''
    对landmark训练图像进行数据增强
    :param image: np.array()类型
    :param landmark_label: 列表：landmark
    :return:
    '''
    image = Image.fromarray(image)
    landmark = np.array(landmark, dtype=np.float)

    image, box = randomCrop_pil(image)
    landmark = corp_lm(landmark, box)
    image, flip_ = randomFlip_pil(image)
    if flip_:
        landmark = getMirrorLM(landmark, image.size[0])
        box = getMirrorBox(box, image.size[0])
    image = randomColor_pil(image)
    image, landmark = randomRotation_pil(image, landmark)
    image_w, image_h = image.size
    box[2] = image_w
    box[3] = image_h
    return np.array(image), landmark, box

def Rotation_point(points, angle, image_w, image_h):
    angle = angle / 180.0 * PI
    points = points - [image_w / 2, image_h / 2]

    rotat = [[np.cos(angle), -np.sin(angle)], [np.sin(angle), np.cos(angle)]]
    rotat = np.array(rotat)
    points = points @ rotat
    points = points + [image_w / 2, image_h / 2]
    return points

def randomRotation_pil(image, lm, angle=30, mode=Image.BICUBIC):
    """
     对图像进行随机任意角度(0~360度)旋转
    :param mode 邻近插值,双线性插值,双三次B样条插值(default)
    :param image PIL的图像image
    :return: 旋转转之后的图像
    """
    rate = np.random.randint(0, 100)
    if rate > 50:
        return image, lm

    random_angle = np.random.randint(-angle, angle)

    lm = Rotation_point(lm, random_angle, image.size[0], image.size[1])

    return image.rotate(random_angle, mode), lm

def randomFlip_pil(image):
    '''
    #图像翻转（类似于镜像，镜子中的自己）
    #FLIP_LEFT_RIGHT,左右翻转
    #FLIP_TOP_BOTTOM,上下翻转
    #ROTATE_90, ROTATE_180, or ROTATE_270.按照角度进行旋转，与randomRotate()功能类似
    :param image:
    :return:
    '''
    rate = np.random.randint(0, 100)
    if rate > 50:
        return image.transpose(Image.FLIP_LEFT_RIGHT), True
    else:
        return image, False

def randomCrop_pil(image):
    """
    对图像随意剪切,考虑到图像大小范围(68,68),使用一个一个大于(36*36)的窗口进行截图
    :param image: PIL的图像image
    :return: 剪切之后的图像
    """
    image_width = image.size[0]
    image_height = image.size[1]
    rate = np.random.randint(0,100)
    if rate > 50:
        return image, [0, 0, image_width, image_height]

    new_w = np.random.randint(int(image_width * 0.85), image_width)
    new_h = np.random.randint(int(image_height * 0.85), image_height)
    new_x = np.random.randint(0, image_width - new_w)
    new_y = np.random.randint(0, image_height - new_h)
    x_2 = new_x + new_w
    y_2 = new_y + new_h

    return image.crop([new_x, new_y, x_2, y_2]), [new_x, new_y, new_w, new_h]

def randomColor_pil(image):
    """
    对图像进行颜色抖动
    :param image: PIL的图像image
    :return: 有颜色色差的图像image
    """
    random_factor = np.random.randint(0, 25) / 10.  # 随机因子
    color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
    random_factor = np.random.randint(5, 15) / 10.  # 随机因子
    brightness_image = ImageEnhance.Brightness(color_image).enhance(random_factor)  # 调整图像的亮度
    random_factor = np.random.randint(10, 21) / 10.  # 随机因1子
    contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
    random_factor = np.random.randint(0, 25) / 10.  # 随机因子
    return ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度

def randomGaussian_pil(image, mean=0.2, sigma=0.3):
    """
     对图像进行高斯噪声处理
    :param image:
    :return:
    """

    def gaussianNoisy(im, mean=0.2, sigma=0.3):
        """
        对图像做高斯噪音处理
        :param im: 单通道图像
        :param mean: 偏移量
        :param sigma: 标准差
        :return:
        """
        for _i in range(len(im)):
            im[_i] += random.gauss(mean, sigma)
        return im

    # 将图像转化成数组
    img = np.asarray(image)
    img.flags.writeable = True  # 将数组改为读写模式
    width, height = img.shape[:2]
    img_r = gaussianNoisy(img[:, :, 0].flatten(), mean, sigma)
    img_g = gaussianNoisy(img[:, :, 1].flatten(), mean, sigma)
    img_b = gaussianNoisy(img[:, :, 2].flatten(), mean, sigma)
    img[:, :, 0] = img_r.reshape([width, height])
    img[:, :, 1] = img_g.reshape([width, height])
    img[:, :, 2] = img_b.reshape([width, height])
    return Image.fromarray(np.uint8(img))

def corp_lm(lm, bbox):
    if len(bbox) != 4:
        return []
    lm = lm.reshape(-1, 2)
    lm = lm - bbox[:2]
    return lm

def test_rotation(imgList, width=800, heigth=800, badCasePath=None, show=True, save=False):
    cv2.namedWindow("ShowImage", 0)
    cv2.moveWindow("ShowImage", 100, 100)
    cv2.resizeWindow("ShowImage", width,heigth)
    id, adds = 0, 1

    while id < len(imgList):
        imgfile = os.path.splitext(imgList[id])[0] + ".png"
        if not checkFileExist(imgfile):
            id = id + adds
            continue
        print(f"{id}/{len(imgList)}: {imgfile}")
        image = Image.open(imgfile)
        bbox = readBBox(imgfile)
        bbox = bbox.bbox
        image = image.crop(XYWH_2_XYXY(bbox))
        img_src = np.array(image, dtype=np.uint8)

        landmark_gt_file = os.path.splitext(imgfile)[0] + ".gt17"
        landmark_gt = readPoint2dFromFile(landmark_gt_file)
        landmark_gt = np.array(landmark_gt, dtype=np.float)
        landmark_gt = landmark_gt - bbox[:2]

        print(image.size)
        image = np.array(image)
        print(image.shape)
        image, landmark = lm_image_enhance(image, landmark_gt)

        image = drawPoints(image, landmark.tolist(), color=(0, 255, 0))
        # img_src = drawPoints(img_src, landmark_gt, color=(0, 255, 0))

        # image = cv2.circle(image, (100, 100), 2, (255, 255, 255))
        # image = cv2.circle(image, (100, 200), 2, (255, 255, 255))

        if show:
            id, adds, save = showOperates(image, "ShowImage", id, adds)
            # id, adds, save = showOperates(img_src, "ShowImage", id, adds)
        else:
            id += 1

        # if save and badCasePath != None:
        #     saveImage(img, image, imgfile, badCasePath)

def test_flip(imgList, width=800, heigth=800, badCasePath=None, show=True, save=False):
    cv2.namedWindow("ShowImage", 0)
    cv2.moveWindow("ShowImage", 100, 100)
    cv2.resizeWindow("ShowImage", width, heigth)
    id, adds = 0, 1

    while id < len(imgList):
        imgfile = os.path.splitext(imgList[id])[0] + ".png"
        if not checkFileExist(imgfile):
            id = id + adds
            continue
        print(f"{id}/{len(imgList)}: {imgfile}")
        image = Image.open(imgfile)
        bbox = readBBox(imgfile)
        bbox = bbox.bbox
        bbox = XYWH_2_XYXY(bbox)
        image = image.crop(bbox)
        img_src = np.array(image, dtype=np.uint8)

        # image = randomFlip_pil(image)
        image = randomGaussian_pil(image)
        image = np.array(image, dtype=np.uint8)

        if show:
            id, adds, save = showOperates(image, "ShowImage", id, adds)
            # id, adds, save = showOperates(img_src, "ShowImage", id, adds)
        else:
            id += 1

if __name__ == '__main__':
    imagelist = [
        '/nas/untouch_data/TrainData/industrial_camera_gray/outdoor/dahua/0208-0216/all_images_labelme.list',
    ]
    imglist = readFileList(imagelist[0])
    test_rotation(imglist)